import argparse
import bson
import time
from datetime import datetime
import logging
import sys

from pymongo import MongoClient

logging.basicConfig(
    stream=sys.stdout, level=logging.INFO, format="%(asctime)s-%(levelname)s-%(message)s"
)


class Worker(object):
    intervalWords = {
        "H": "时线", "D": "日线", "W": "周线", "M": "月线"
    }
    marketChoices = {
        "stock": {
            "D": bson.ObjectId("6695cf84206a1589b0423400"),
            "W": bson.ObjectId("6683934bc31460f5fdedf38f"),
            "M": bson.ObjectId("66839372c31460f5fdedf390"),
        }, "future": {
            "H": bson.ObjectId("66716adc53a86f32e752c4cf"),
            "D": bson.ObjectId("667926c393d69c87191c3655"),
            "W": bson.ObjectId("6614a8f08b2da599649ceea4"),
            "M": bson.ObjectId("66794d1493d69c87191c76fd")
        }, "forex": {
            "H": bson.ObjectId("668392b9c31460f5fdedf38d"),
            "D": bson.ObjectId("65a69549f890ca1e6f373ea5"),
            "W": bson.ObjectId("6614a9708b2da599649ceea6"),
            "M": bson.ObjectId("6614a98a8b2da599649ceea7")
        }, "digiccy": {
            "H": bson.ObjectId("65a695b0f890ca1e6f373ea8"),
            "D": bson.ObjectId("66794dac93d69c87191c76fe"),
            "W": bson.ObjectId("66794de393d69c87191c76ff"),
            "M": bson.ObjectId("66794e0793d69c87191c7700")
        }}

    def __init__(self, url):
        self.client = MongoClient(url)

    def build_strategy(self, user, market, interval, oid):
        owned = self.client.crystal.strategy.count_documents({"user": user, "market": market, "interval": interval})
        if owned > 0:
            logging.info(f"拥有策略{owned}，跳过")
            return None
        document = self.client.crystal.strategy.find_one({"_id": oid})
        if document is None:
            logging.info(f"没有策略{oid}，跳过")
            return None
        result = self.client.crystal.strategy.insert_one({
            "user": user, "market": market, "interval": interval, "title": document["title"],
            "config": document["config"], "params": document["params"], "locked": True, "state": 1,
            "sequence": 0, "created": datetime.now(), "updated": datetime.now()
        })
        return result.inserted_id

    def build_account(self, user, market, strategy, interval):
        result = self.client.crystal.account.insert_one({
            "user": user, "market": market, "strategy": strategy,
            "title": f"{self.intervalWords[interval]}测试", "amount": 100000,
            "fund": 1000000, "cash": 1000000, "available": 1000000, "frozen": 0,
            "locked": True, "state": 0,
            "created": datetime.now(), "updated": datetime.now()
        })
        return result.inserted_id

    def start(self):
        for user in self.client.vauth.user.find():
            logging.info(f"用户{user['_id']}: {user.get('realname')}")
            for market, configs in self.marketChoices.items():
                for interval, oid in configs.items():
                    logging.info(f"{market}-{interval}")
                    strategy = self.build_strategy(user["_id"], market, interval, oid)
                    if strategy is not None:
                        logging.info(f"创建策略{strategy}成功")
                        resp = self.build_account(user["_id"], market, strategy, interval)
                        logging.info(f"创建账户{resp}成功")
        self.client.close()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--url", type=str, default="mongodb://admin:8VvvQ8SpHvctIR2G@106.15.54.243:27017/admin?authSource=admin")
    params, _ = parser.parse_known_args()
    # Worker(params.url).start()
    while True:
        Worker(params.url).start()
        time.sleep(60)


if __name__ == "__main__":
    main()
